import json

from hand.track.umetrack.datasets.video_pose_data import YVRImagePoseStream, load_hand_model_from_dict, \
    YVRDetectImageStream
from .base import Loader


class YVRGTLoader(Loader):
    def __init__(self, data_path: str):
        self._image_stream = YVRImagePoseStream(data_path, time_order=True)
        with open(r"C:\workspace\hand_xr\data\umetrack\example\generic_hand_model.json") as f:
            self.hand_model = load_hand_model_from_dict(json.load(f))

    def __len__(self) -> int:
        return len(self._image_stream)

    def __iter__(self):
        for input_frame, gt_tracking in self._image_stream:
            yield input_frame, gt_tracking


class YVRDetLoader(Loader):
    def __init__(self, data_path: str):
        self._image_stream = YVRDetectImageStream(data_path)
        with open(r"C:\workspace\hand_xr\data\umetrack\example\generic_hand_model.json") as f:
            self.hand_model = load_hand_model_from_dict(json.load(f))

    def __len__(self) -> int:
        return len(self._image_stream)

    def __iter__(self):
        for input_frame, gt_tracking in self._image_stream:
            yield input_frame, gt_tracking
